# model/model.py
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:, :x.size(1), :]

class Transformer(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, embedding_dim, hidden_dim, num_heads, num_layers, dropout, max_seq_length):
        super(Transformer, self).__init__()
        self.src_vocab_size = src_vocab_size
        self.tgt_vocab_size = tgt_vocab_size
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        
        self.encoder_embedding = nn.Embedding(src_vocab_size, embedding_dim)
        self.decoder_embedding = nn.Embedding(tgt_vocab_size, embedding_dim)
        self.positional_encoding = PositionalEncoding(embedding_dim, max_seq_length)
        
        self.encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim, 
            nhead=num_heads, 
            dim_feedforward=hidden_dim * 4,  # 通常是隐藏层维度的4倍
            dropout=dropout,
            batch_first=True
        )
        self.encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers)
        
        self.decoder_layer = nn.TransformerDecoderLayer(
            d_model=hidden_dim, 
            nhead=num_heads, 
            dim_feedforward=hidden_dim * 4,  # 通常是隐藏层维度的4倍
            dropout=dropout, 
            batch_first=True
        )
        self.decoder = nn.TransformerDecoder(self.decoder_layer, num_layers=num_layers)
        
        self.fc_out = nn.Linear(hidden_dim, tgt_vocab_size)
        
        # 初始化权重
        self._reset_parameters()
    
    def _reset_parameters(self):
        """初始化模型参数"""
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)
    
    def forward(self, src, tgt, src_mask, tgt_mask):
        # 源序列和目标序列的嵌入
        src_embedded = self.positional_encoding(self.encoder_embedding(src))
        tgt_embedded = self.positional_encoding(self.decoder_embedding(tgt))
        
        # 创建注意力掩码
        # src_mask 和 tgt_mask 是 padding mask，形状为 [batch_size, seq_len]
        # 需要转换为 PyTorch Transformer 所需的格式
        
        # 为解码器创建因果掩码（确保只能看到当前及之前的位置）
        tgt_len = tgt.size(1)
        tgt_attn_mask = self._generate_square_subsequent_mask(tgt_len).to(tgt.device)
        
        # 编码器前向传播
        # 注意：src_mask 需要反转，因为 PyTorch Transformer 中 1 表示要遮蔽的位置
        memory = self.encoder(src_embedded, src_key_padding_mask=src_mask)
        
        # 解码器前向传播
        # 使用 tgt_attn_mask 作为 tgt_mask 参数，tgt_mask 作为 tgt_key_padding_mask
        output = self.decoder(
            tgt_embedded, 
            memory, 
            tgt_mask=tgt_attn_mask,
                                     memory_key_padding_mask=src_mask,
            tgt_key_padding_mask=tgt_mask
        )
        
        # 输出层
        output = self.fc_out(output)
        return output
        
    def _generate_square_subsequent_mask(self, sz):
        """生成用于解码器的方形后续掩码（确保每个位置只能关注之前的位置）"""
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask
        
    def encode(self, src, src_mask):
        """单独的编码器函数，用于推理"""
        src_embedded = self.positional_encoding(self.encoder_embedding(src))
        return self.encoder(src_embedded, src_key_padding_mask=src_mask)
        
    def decode(self, tgt, memory, tgt_mask, src_mask):
        """单独的解码器函数，用于推理"""
        tgt_embedded = self.positional_encoding(self.decoder_embedding(tgt))
        tgt_len = tgt.size(1)
        tgt_attn_mask = self._generate_square_subsequent_mask(tgt_len).to(tgt.device)
        
        return self.decoder(
            tgt_embedded, 
            memory, 
            tgt_mask=tgt_attn_mask,
            memory_key_padding_mask=src_mask,
            tgt_key_padding_mask=tgt_mask
        )